import java.util.*;

/**
 * @Author: Liu, Yunfeng
 * @DateTime: 2021/7/28 21:27
 * @Description: TODO
 */
public class distanceK {
    Map<Integer, TreeNode> parent = new HashMap<>();
    List<Integer> ans = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        getAll(root);
        getAns(target, target, k, 0);
        return ans;
    }

    public void getAll(TreeNode root) {
        if (root.left != null) {
            parent.put(root.left.val, root);
            getAll(root.left);
        }
        if (root.right != null) {
            parent.put(root.right.val, root);
            getAll(root.right);
        }
    }

    public void getAns(TreeNode current, TreeNode from, int k, int now) {
        if (current == null) {
            return;
        }
        if (now == k) {
            ans.add(current.val);
            return;
        }
        if (current.left != from) {
            getAns(current.left, current, k, now + 1);
        }
        if (current.right != from) {
            getAns(current.right, current, k, now + 1);
        }
        TreeNode parentNode = parent.get(current.val);
        if (parentNode != from) {
            getAns(parentNode, current, k, now + 1);
        }
    }
}
